import logging
import asyncio
import json
from typing import Dict, Any
from celery import current_task
from sqlalchemy.orm import Session

from app.celery_app import celery_app
from app.db.session import SessionLocal
from app.crud.crud_generation_job import update_job
from app.models.generation_job import GenerationJob, JobStatus
from app.services.ai_manager import ai_service_manager
from app.services.storage_manager import storage_manager
from app.services.ai_base import GenerationRequest, AIServiceType, AIServiceProvider
from app.services.storage_base import FileType
from app.core.config import settings
from app.api.v1.video_service import update_quota_usage

logger = logging.getLogger(__name__)


@celery_app.task(bind=True)
def generate_image_task(self, job_id: int) -> Dict[str, Any]:
    """
    异步处理图像生成任务
    """
    db: Session = SessionLocal()
    
    try:
        # 更新任务状态为处理中
        job = db.query(GenerationJob).filter(GenerationJob.id == job_id).first()
        if not job:
            raise ValueError(f"Job {job_id} not found")
            
        update_job(db, db_obj=job, obj_in={"status": JobStatus.PROCESSING})
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 10, "total": 100, "status": "开始生成图片..."})
        
        # 获取AI服务
        ai_service = ai_service_manager.get_default_service()
        if not ai_service:
            raise RuntimeError("AI服务不可用")
        
        # 创建AI请求
        ai_request = GenerationRequest(
            service_type=AIServiceType.IMAGE_GENERATION,
            prompt=job.prompt,
            params=job.params or {},
            user_id=job.user_id
        )
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 30, "total": 100, "status": "正在生成图片..."})
        
        # 调用AI服务生成图片
        ai_response = asyncio.run(ai_service.generate_image(ai_request))
        
        if not ai_response.success:
            raise RuntimeError(f"图片生成失败: {ai_response.error_message}")
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 70, "total": 100, "status": "图片生成完成，上传到存储..."})
        
        # 获取存储服务并保存图片
        storage_service = storage_manager.get_service()
        if not storage_service:
            raise RuntimeError("存储服务不可用")
        
        # 保存生成的图片
        storage_result = asyncio.run(storage_service.upload_file(
            file_data=ai_response.content,
            file_name=f"generated_image_{job_id}.png",
            file_type=FileType.GENERATED_IMAGE,
            content_type="image/png"
        ))
        
        if not storage_result.success:
            raise RuntimeError(f"图片保存失败: {storage_result.error_message}")
        
        result_url = storage_result.file_url
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 90, "total": 100, "status": "保存完成，更新数据库..."})
        
        # 更新任务为完成状态
        update_job(db, db_obj=job, obj_in={
            "status": JobStatus.COMPLETED,
            "result_url": result_url
        })
        
        self.update_state(state="SUCCESS", meta={"current": 100, "total": 100, "status": "任务完成"})
        
        return {
            "job_id": job_id,
            "status": "completed",
            "result_url": result_url
        }
        
    except Exception as exc:
        logger.error(f"图像生成任务失败: {exc}")
        
        # 更新任务为失败状态
        if 'job' in locals():
            update_job(db, db_obj=job, obj_in={"status": JobStatus.FAILED})
            
        self.update_state(
            state="FAILURE",
            meta={"current": 100, "total": 100, "status": f"任务失败: {str(exc)}"}
        )
        raise exc
        
    finally:
        db.close()


@celery_app.task(bind=True)
def generate_video_task(self, job_id: int) -> Dict[str, Any]:
    """
    异步处理视频生成任务
    """
    db: Session = SessionLocal()
    
    try:
        # 更新任务状态为处理中
        job = db.query(GenerationJob).filter(GenerationJob.id == job_id).first()
        if not job:
            raise ValueError(f"Job {job_id} not found")
            
        update_job(db, db_obj=job, obj_in={"status": JobStatus.PROCESSING})
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 10, "total": 100, "status": "开始生成视频..."})
        
        # 获取AI服务
        ai_service = ai_service_manager.get_default_service()
        if not ai_service:
            raise RuntimeError("AI服务不可用")
        
        # 创建AI请求
        ai_request = GenerationRequest(
            service_type=AIServiceType.VIDEO_GENERATION,
            prompt=job.prompt,
            params=job.params or {},
            user_id=job.user_id
        )
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 30, "total": 100, "status": "正在生成视频..."})
        
        # 调用AI服务生成视频
        ai_response = asyncio.run(ai_service.generate_video(ai_request))
        
        if not ai_response.success:
            raise RuntimeError(f"视频生成失败: {ai_response.error_message}")
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 70, "total": 100, "status": "视频生成完成，上传到存储..."})
        
        # 获取存储服务并保存视频
        storage_service = storage_manager.get_service()
        if not storage_service:
            raise RuntimeError("存储服务不可用")
        
        # 保存生成的视频
        storage_result = asyncio.run(storage_service.upload_file(
            file_data=ai_response.content,
            file_name=f"generated_video_{job_id}.mp4",
            file_type=FileType.GENERATED_VIDEO,
            content_type="video/mp4"
        ))
        
        if not storage_result.success:
            raise RuntimeError(f"视频保存失败: {storage_result.error_message}")
        
        result_url = storage_result.file_url
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={"current": 90, "total": 100, "status": "保存完成，更新数据库..."})
        
        # 更新任务为完成状态
        update_job(db, db_obj=job, obj_in={
            "status": JobStatus.COMPLETED,
            "result_url": result_url
        })
        
        self.update_state(state="SUCCESS", meta={"current": 100, "total": 100, "status": "任务完成"})
        
        return {
            "job_id": job_id,
            "status": "completed",
            "result_url": result_url
        }
        
    except Exception as exc:
        logger.error(f"视频生成任务失败: {exc}")
        
        # 更新任务为失败状态
        if 'job' in locals():
            update_job(db, db_obj=job, obj_in={"status": JobStatus.FAILED})
            
        self.update_state(
            state="FAILURE",
            meta={"current": 100, "total": 100, "status": f"任务失败: {str(exc)}"}
        )
        raise exc
        
    finally:
        db.close()


@celery_app.task(bind=True)
def image_edit_task(self, job_id: int, source_image_url: str) -> Dict[str, Any]:
    """
    异步处理图像编辑任务
    """
    db: Session = SessionLocal()
    
    try:
        job = db.query(GenerationJob).filter(GenerationJob.id == job_id).first()
        if not job:
            raise ValueError(f"Job {job_id} not found")
            
        update_job(db, db_obj=job, obj_in={"status": JobStatus.PROCESSING})
        
        self.update_state(state="PROGRESS", meta={"current": 10, "total": 100, "status": "开始编辑图片..."})
        
        # 获取AI服务
        ai_service = ai_service_manager.get_default_service()
        if not ai_service:
            raise RuntimeError("AI服务不可用")
        
        # 创建AI请求
        ai_request = GenerationRequest(
            service_type=AIServiceType.IMAGE_EDITING,
            prompt=job.prompt,
            params={
                "source_image_url": source_image_url,
                **(job.params or {})
            },
            user_id=job.user_id
        )
        
        self.update_state(state="PROGRESS", meta={"current": 30, "total": 100, "status": "正在编辑图片..."})
        
        # 调用AI服务编辑图片
        ai_response = asyncio.run(ai_service.edit_image(ai_request))
        
        if not ai_response.success:
            raise RuntimeError(f"图片编辑失败: {ai_response.error_message}")
        
        self.update_state(state="PROGRESS", meta={"current": 70, "total": 100, "status": "编辑完成，上传到存储..."})
        
        # 获取存储服务并保存图片
        storage_service = storage_manager.get_service()
        if not storage_service:
            raise RuntimeError("存储服务不可用")
        
        storage_result = asyncio.run(storage_service.upload_file(
            file_data=ai_response.content,
            file_name=f"edited_image_{job_id}.png",
            file_type=FileType.GENERATED_IMAGE,
            content_type="image/png"
        ))
        
        if not storage_result.success:
            raise RuntimeError(f"图片保存失败: {storage_result.error_message}")
        
        result_url = storage_result.file_url
        
        self.update_state(state="PROGRESS", meta={"current": 90, "total": 100, "status": "保存完成，更新数据库..."})
        
        # 更新任务为完成状态
        update_job(db, db_obj=job, obj_in={
            "status": JobStatus.COMPLETED,
            "result_url": result_url
        })
        
        self.update_state(state="SUCCESS", meta={"current": 100, "total": 100, "status": "任务完成"})
        
        return {
            "job_id": job_id,
            "status": "completed",
            "result_url": result_url
        }
        
    except Exception as exc:
        logger.error(f"图像编辑任务失败: {exc}")
        
        if 'job' in locals():
            update_job(db, db_obj=job, obj_in={"status": JobStatus.FAILED})
            
        self.update_state(
            state="FAILURE",
            meta={"current": 100, "total": 100, "status": f"任务失败: {str(exc)}"}
        )
        raise exc
        
    finally:
        db.close()


@celery_app.task(bind=True)
def face_swap_task(self, job_id: int, source_image_url: str, target_face_url: str) -> Dict[str, Any]:
    """
    异步处理人脸替换任务
    """
    db: Session = SessionLocal()
    
    try:
        job = db.query(GenerationJob).filter(GenerationJob.id == job_id).first()
        if not job:
            raise ValueError(f"Job {job_id} not found")
            
        update_job(db, db_obj=job, obj_in={"status": JobStatus.PROCESSING})
        
        self.update_state(state="PROGRESS", meta={"current": 10, "total": 100, "status": "开始人脸替换..."})
        
        # 获取AI服务
        ai_service = ai_service_manager.get_default_service()
        if not ai_service:
            raise RuntimeError("AI服务不可用")
        
        # 创建AI请求
        ai_request = GenerationRequest(
            service_type=AIServiceType.FACE_SWAP,
            prompt=job.prompt or "人脸替换",
            params={
                "source_image_url": source_image_url,
                "target_face_url": target_face_url,
                **(job.params or {})
            },
            user_id=job.user_id
        )
        
        self.update_state(state="PROGRESS", meta={"current": 30, "total": 100, "status": "正在处理人脸替换..."})
        
        # 调用AI服务进行人脸替换
        ai_response = asyncio.run(ai_service.face_swap(ai_request))
        
        if not ai_response.success:
            raise RuntimeError(f"人脸替换失败: {ai_response.error_message}")
        
        self.update_state(state="PROGRESS", meta={"current": 70, "total": 100, "status": "替换完成，上传到存储..."})
        
        # 获取存储服务并保存图片
        storage_service = storage_manager.get_service()
        if not storage_service:
            raise RuntimeError("存储服务不可用")
        
        storage_result = asyncio.run(storage_service.upload_file(
            file_data=ai_response.content,
            file_name=f"face_swap_{job_id}.png",
            file_type=FileType.GENERATED_IMAGE,
            content_type="image/png"
        ))
        
        if not storage_result.success:
            raise RuntimeError(f"图片保存失败: {storage_result.error_message}")
        
        result_url = storage_result.file_url
        
        self.update_state(state="PROGRESS", meta={"current": 90, "total": 100, "status": "保存完成，更新数据库..."})
        
        # 更新任务为完成状态
        update_job(db, db_obj=job, obj_in={
            "status": JobStatus.COMPLETED,
            "result_url": result_url
        })
        
        self.update_state(state="SUCCESS", meta={"current": 100, "total": 100, "status": "任务完成"})
        
        return {
            "job_id": job_id,
            "status": "completed",
            "result_url": result_url
        }
        
    except Exception as exc:
        logger.error(f"人脸替换任务失败: {exc}")
        
        if 'job' in locals():
            update_job(db, db_obj=job, obj_in={"status": JobStatus.FAILED})
            
        self.update_state(
            state="FAILURE",
            meta={"current": 100, "total": 100, "status": f"任务失败: {str(exc)}"}
        )
        raise exc
        
    finally:
        db.close()

@celery_app.task(bind=True, queue='jimeng_video')
def generate_video_async(self, request_data: Dict[str, Any], user_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
    """
    异步视频生成任务 - 专为即梦AI设计
    使用专用队列确保并发限制=1
    """
    try:
        logger.info(f"开始处理异步视频生成任务 - 用户: {user_id}, 任务ID: {self.request.id}")
        
        # 更新任务状态为处理中
        self.update_state(state="PROGRESS", meta={
            "current": 10, 
            "total": 100, 
            "status": "开始生成视频...",
            "step": "initializing"
        })
        
        # 确保AI服务管理器已初始化
        ai_service_manager.register_all_enabled_services()
        
        # 构建AI服务请求
        ai_request = GenerationRequest(
            prompt=request_data.get("prompt"),
            service_type=AIServiceType.VIDEO_GENERATION,
            params=params,
            user_id=user_id
        )
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={
            "current": 30, 
            "total": 100, 
            "status": "正在调用即梦AI生成视频...",
            "step": "generating"
        })
        
        # 获取即梦AI服务
        service = ai_service_manager.get_service(AIServiceProvider.JIMENG_AI)
        if not service:
            raise RuntimeError("即梦AI服务初始化失败，请检查配置")
        
        logger.info(f"使用即梦AI服务生成视频 - 提示词: {ai_request.prompt}")
        logger.info(f"📋 生成参数详情: {json.dumps(params, ensure_ascii=False)}")
        
        # 调用AI服务生成视频（这里会等待5分钟）
        logger.info(f"🚀 开始调用即梦AI 720p接口...")
        result = asyncio.run(service.generate_video(ai_request))
        logger.info(f"📥 即梦AI调用完成，成功: {result.success}")
        
        if result.success:
            logger.info(f"✅ 视频生成成功!")
            logger.info(f"   视频URL: {result.result_url}")
            logger.info(f"   成本: {result.cost}")
            logger.info(f"   元数据: {json.dumps(result.metadata or {}, ensure_ascii=False)}")
        else:
            logger.error(f"❌ 视频生成失败: {result.error_message}")
            if hasattr(result, 'metadata') and result.metadata:
                logger.error(f"   错误元数据: {json.dumps(result.metadata, ensure_ascii=False)}")
        
        if not result.success:
            raise RuntimeError(f"视频生成失败: {result.error_message}")
        
        # 更新任务进度
        self.update_state(state="PROGRESS", meta={
            "current": 80, 
            "total": 100, 
            "status": "视频生成完成，正在更新配额...",
            "step": "finalizing"
        })
        
        # 更新用户配额
        asyncio.run(update_quota_usage(user_id, params, result.cost or 0.0))
        
        # 准备完整的任务结果
        task_result = {
            "success": True,
            "task_id": result.metadata.get("task_id") if result.metadata else self.request.id,
            "video_url": result.result_url,
            "cost": result.cost,
            "generation_type": params.get("generation_type", "text_to_video"),
            "aspect_ratio": request_data.get("aspect_ratio"),
            "duration": request_data.get("duration"),
            "metadata": result.metadata or {}
        }
        
        # 更新任务进度，包含完整结果信息
        self.update_state(state="SUCCESS", meta={
            "current": 100, 
            "total": 100, 
            "status": "任务完成",
            "step": "completed",
            # 包含完整的任务结果供API查询使用
            "video_url": result.result_url,
            "cost": result.cost,
            "generation_type": params.get("generation_type", "text_to_video"),
            "aspect_ratio": request_data.get("aspect_ratio"),
            "duration": request_data.get("duration"),
            "metadata": result.metadata or {}
        })
        
        logger.info(f"视频生成任务完成 - 任务ID: {self.request.id}, URL: {result.result_url}")
        
        return task_result
        
    except Exception as exc:
        logger.error(f"异步视频生成任务失败: {exc}")
        
        # 修复：直接更新状态并返回结果，不要再次序列化异常
        error_result = {
            "success": False,
            "error_message": str(exc),
            "task_id": self.request.id,
            "generation_type": params.get("generation_type", "text_to_video"),
            "aspect_ratio": request_data.get("aspect_ratio"),
            "duration": request_data.get("duration")
        }
        
        # 直接更新状态为SUCCESS并返回错误结果
        self.update_state(
            state="SUCCESS",  # 改为SUCCESS避免异常序列化
            meta={
                "current": 100, 
                "total": 100, 
                "status": f"任务失败: {str(exc)}",
                "step": "failed",
                "error": str(exc)
            }
        )
        
        # 返回错误结果而不是抛出异常
        return error_result
